package tcpserver

import (
	"encoding/binary"
	"io"
	"os"

	"github.com/juju/ratelimit"
)

/////////////////////////////////////
//  写入协议头
func WriteHeader(w io.Writer, length uint32, Series uint16) {

	var header DataHeader

	header.Reset()
	header.DataLength = length
	header.Series = Series
	binary.Write(w, binary.LittleEndian, &header)
}

// 发送数据尾部
func WriteFooter(w io.Writer) {

	var d DataFooter

	d.Reset()
	binary.Write(w, binary.LittleEndian, &d)
}

// 发送文件数据到远程
func WriteFile(w io.Writer, filepath string, offset int64, length int64) (int64, error) {

	f, err := os.Open(filepath)
	if err != nil {
		return 0, err
	}

	// 退出函数后关闭文件
	defer f.Close()

	// 检查文件大小
	fi, err2 := f.Stat()
	if err2 != nil {
		return 0, err2
	}

	offset, length = CalcuteTransferParam(fi.Size(), offset, length)

	// 如果文件大小为0, 直接返回
	if length <= 0 {
		return 0, nil
	}

	// 移动到指定的位置
	if offset > 0 {
		f.Seek(offset, 0)
	}

	// 发送指定的大小的数据
	return io.CopyN(w, f, length)
}

// 根据速度来发送数据
func WriteFileAtSpeed(w io.Writer, filepath string, offset int64, length int64, bytesPerSecond int64) (int64, error) {

	f, err := os.Open(filepath)
	if err != nil {
		return 0, err
	}

	// 退出函数后关闭文件
	defer f.Close()

	// 检查文件大小
	fi, err2 := f.Stat()
	if err2 != nil {
		return 0, err2
	}

	offset, length = CalcuteTransferParam(fi.Size(), offset, length)

	// 如果文件大小为0, 直接返回
	if length <= 0 {
		return 0, nil
	}

	// 移动到指定的位置
	if offset > 0 {
		f.Seek(offset, 0)
	}

	// 限速
	// Source holding 1MB
	//src := bytes.NewReader(make([]byte, bytesPerSecond))

	// Bucket adding 100KB every second, holding max 100KB
	//bucket := ratelimit.NewBucketWithRate(100*1024, 100*1024)
	bucket := ratelimit.NewBucketWithRate(float64(bytesPerSecond), bytesPerSecond)

	// 发送指定的大小的数据
	return io.CopyN(w, ratelimit.Reader(f, bucket), length)
}

// 计算传输的偏移和大小
func CalcuteTransferParam(total int64, offset int64, length int64) (newOffset int64, newLength int64) {
	newOffset = offset
	newLength = length

	if newOffset < 0 || newOffset >= total {
		newOffset = 0
	}

	if newLength <= 0 {
		newLength = total
	}

	if newOffset+newLength > total {
		newLength = total - newOffset
	}

	return
}
